#include "param.h"
#include "types.h"
#include "defs.h"
#include "memlayout.h"
#include "riscv.h"

extern char etext[]; // kernel.ld sets this to end of kernel code.
extern char trampoline[]; // trampoline.S

// Make a direct-map page table for the kernel.
pagetable_t kernel_pagetable;

pte_t* walk(pagetable_t, uint64, int);

// 便捷封装：将 va..va+sz-1 映射到 pa..pa+sz-1，出错直接 panic
static void
kvmmap(uint64 va, uint64 pa, uint64 sz, int perm)
{
  if(mappages(kernel_pagetable, va, pa, sz, perm) != 0)
    panic("kvmmap");
}

// 为内核创建一个直接映射的页表
void
kvminit(void)
{
  printf("kvminit: 创建内核页表...\n"); // 检查点
  kernel_pagetable = (pagetable_t) kalloc();
  memset(kernel_pagetable, 0, PGSIZE);
  printf("kvminit: 根页表已分配在 %p\n", kernel_pagetable); //检查点

  // 映射 uart 寄存器（保持你原有写法也可）
  printf("kvminit: 映射 UART...\n");
  mappages(kernel_pagetable, UART0, UART0, PGSIZE, PTE_R | PTE_W);

  // 映射 VIRTIO 设备
  printf("kvminit: 映射 VIRTIO...\n");
  kvmmap(VIRTIO0, VIRTIO0, PGSIZE, PTE_R | PTE_W);

  // 映射 CLINT（mtime/mtimecmp 在此区间）
  printf("kvminit: 映射 CLINT...\n");
  kvmmap(0x02000000L, 0x02000000L, 0x10000, PTE_R | PTE_W);

  // 映射内核代码段 (可读/可执行)
  printf("kvminit: 映射内核代码段...\n");
  mappages(kernel_pagetable, KERNBASE, KERNBASE, (uint64)etext-KERNBASE, PTE_R | PTE_X);

  // 映射内核数据段和剩余的物理内存 (可读/可写)
  printf("kvminit: 映射内核数据段和物理内存...\n");
  mappages(kernel_pagetable, (uint64)etext, (uint64)etext, PHYSTOP-(uint64)etext, PTE_R | PTE_W);
}


// 将硬件页表寄存器切换到内核页表, 并启用分页
void
kvminithart()
{
  printf("kvminithart: 启用分页机制...\n"); // 检查点
  printf("kvminithart: satp = %p\n", MAKE_SATP(kernel_pagetable)); // <-- 新增检查点
  w_satp(MAKE_SATP(kernel_pagetable));
  sfence_vma();
  printf("kvminithart: 分页已启用!\n"); // 检查点
}


// Return the address of the PTE in page table pagetable
// that corresponds to virtual address va.  If alloc!=0,
// create any required page-table pages.
//
// The risc-v Sv39 scheme has three levels of page-table
// pages. A page-table page contains 512 64-bit PTEs.
// A 64-bit virtual address is split into five fields:
//   39..63 -- must be zero.
//   30..38 -- 9 bits of level-2 index.
//   21..29 -- 9 bits of level-1 index.
//   12..20 -- 9 bits of level-0 index.
//    0..11 -- 12 bits of byte offset within the page.
pte_t*
walk(pagetable_t pagetable, uint64 va, int alloc)
{
  if(va >= (1L << 39))
    panic("walk");

  for(int level = 2; level > 0; level--) {
    pte_t *pte = &pagetable[PX(level, va)];
    if(*pte & PTE_V) {
      pagetable = (pagetable_t)PTE2PA(*pte);
    } else {
      if(!alloc || (pagetable = (pagetable_t)kalloc()) == 0)
        return 0;
      memset(pagetable, 0, PGSIZE);
      *pte = PA2PTE(pagetable) | PTE_V;
    }
  }
  return &pagetable[PX(0, va)];
}

// Create PTEs for virtual addresses starting at va that refer to
// physical addresses starting at pa. va and size might not
// be page-aligned. Returns 0 on success, -1 if walk() couldn't
// allocate a needed page-table page.
int
mappages(pagetable_t pagetable, uint64 va, uint64 pa, uint64 size, int perm)
{
  uint64 a, last;
  pte_t *pte;

  a = PGROUNDDOWN(va);
  last = PGROUNDDOWN(va + size - 1);
  for(;;){
    if((pte = walk(pagetable, a, 1)) == 0)
      return -1;
    if(*pte & PTE_V)
      panic("mappages: remap");
    *pte = PA2PTE(pa) | perm | PTE_V;
    if(a == last)
      break;
    a += PGSIZE;
    pa += PGSIZE;
  }
  return 0;
}